import asyncio
import os
import torch
import mmcv
import numpy as np
import matplotlib.pyplot as plt
from argparse import ArgumentParser
from pathlib import Path

from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector)
from mmdet.core.visualization import imshow_det_bboxes


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('img', help='Image file or directory')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='bbox score threshold')
    parser.add_argument(
        '--async-test',
        action='store_true',
        help='whether to set async options for async inference.')
    parser.add_argument(
        '--output-dir', type=str, default='output', help='Output directory for visualization results')
    parser.add_argument(
        '--image-list', type=str, help='Path to a text file containing image names to process (one per line)')
    args = parser.parse_args()
    return args


def process_single_image(model, img_path, output_dir, score_thr):
    """处理单张图片并保存结果"""
    result = inference_detector(model, img_path)
    # 创建输出文件名
    output_path = os.path.join(output_dir, f'{os.path.basename(img_path)}')
    
    # 读取原始图像
    img = mmcv.imread(img_path)
    img = mmcv.bgr2rgb(img)
    
    # 处理检测结果
    if isinstance(result, tuple):
        bbox_result, segm_result = result
        if isinstance(segm_result, tuple):
            segm_result = segm_result[0]
    else:
        bbox_result, segm_result = result, None
    
    # 将检测结果转换为numpy数组
    if isinstance(bbox_result, list):
        bboxes = []
        labels = []
        for i, bbox in enumerate(bbox_result):
            if len(bbox) > 0:
                bboxes.append(bbox)
                labels.append(np.full((len(bbox),), i, dtype=np.int64))
        if len(bboxes) > 0:
            bboxes = np.concatenate(bboxes, axis=0)
            labels = np.concatenate(labels, axis=0)
        else:
            bboxes = np.zeros((0, 5))
            labels = np.zeros((0,), dtype=np.int64)
    else:
        bboxes = bbox_result
        labels = np.zeros((len(bbox_result),), dtype=np.int64)
    
    # 创建新的图像用于可视化
    plt.figure(figsize=(10, 10))
    plt.imshow(img)
    
    # 绘制检测框
    if len(bboxes) > 0:
        for bbox, label in zip(bboxes, labels):
            if bbox[4] >= score_thr:
                x1, y1, x2, y2 = bbox[:4].astype(np.int32)
                plt.gca().add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, 
                                                fill=False, edgecolor='red', linewidth=2))
                plt.text(x1, y1-5, f'{model.CLASSES[label]}: {bbox[4]:.2f}',
                        color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.7))
    
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close()
    
    print(f'处理完成: {img_path} -> {output_path}')


def main(args):
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 检查CUDA设备
    if args.device.startswith('cuda'):
        if not torch.cuda.is_available():
            print('CUDA不可用，使用CPU进行推理')
            args.device = 'cpu'
        else:
            # 确保使用可用的CUDA设备
            device_id = int(args.device.split(':')[1])
            if device_id >= torch.cuda.device_count():
                print(f'警告：请求的CUDA设备{device_id}不可用，使用设备0')
                args.device = 'cuda:0'
    
    # 构建模型
    model = init_detector(args.config, args.checkpoint, device=args.device)
    
    # 处理输入路径
    input_path = Path(args.img)
    if input_path.is_file():
        # 单张图片处理
        process_single_image(model, str(input_path), args.output_dir, args.score_thr)
    elif input_path.is_dir():
        # 获取要处理的图片列表
        image_names = set()
        if args.image_list:
            with open(args.image_list, 'r') as f:
                image_names = {line.strip() for line in f if line.strip()}
            print(f'从文件加载了 {len(image_names)} 个图片名称')
        
        # 处理目录中的所有图片
        image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
        processed_count = 0
        for img_path in input_path.glob('*'):
            if img_path.suffix.lower() in image_extensions:
                # 如果指定了图片列表，只处理列表中的图片
                if args.image_list and img_path.name not in image_names:
                    continue
                process_single_image(model, str(img_path), args.output_dir, args.score_thr)
                processed_count += 1
        
        print(f'总共处理了 {processed_count} 张图片')
    else:
        raise ValueError(f'输入路径不存在: {args.img}')


async def async_main(args):
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 检查CUDA设备
    if args.device.startswith('cuda'):
        if not torch.cuda.is_available():
            print('CUDA不可用，使用CPU进行推理')
            args.device = 'cpu'
        else:
            # 确保使用可用的CUDA设备
            device_id = int(args.device.split(':')[1])
            if device_id >= torch.cuda.device_count():
                print(f'警告：请求的CUDA设备{device_id}不可用，使用设备0')
                args.device = 'cuda:0'
    
    # 构建模型
    model = init_detector(args.config, args.checkpoint, device=args.device)
    
    # 处理输入路径
    input_path = Path(args.img)
    if input_path.is_file():
        # 单张图片处理
        tasks = asyncio.create_task(async_inference_detector(model, str(input_path)))
        result = await asyncio.gather(tasks)
        output_path = os.path.join(args.output_dir, f'{os.path.basename(str(input_path))}')
        
        # 读取原始图像
        img = mmcv.imread(str(input_path))
        img = mmcv.bgr2rgb(img)
        
        # 处理检测结果
        if isinstance(result[0], tuple):
            bbox_result, segm_result = result[0]
            if isinstance(segm_result, tuple):
                segm_result = segm_result[0]
        else:
            bbox_result, segm_result = result[0], None
        
        # 将检测结果转换为numpy数组
        if isinstance(bbox_result, list):
            bboxes = []
            labels = []
            for i, bbox in enumerate(bbox_result):
                if len(bbox) > 0:
                    bboxes.append(bbox)
                    labels.append(np.full((len(bbox),), i, dtype=np.int64))
            if len(bboxes) > 0:
                bboxes = np.concatenate(bboxes, axis=0)
                labels = np.concatenate(labels, axis=0)
            else:
                bboxes = np.zeros((0, 5))
                labels = np.zeros((0,), dtype=np.int64)
        else:
            bboxes = bbox_result
            labels = np.zeros((len(bbox_result),), dtype=np.int64)
        
        # 创建新的图像用于可视化
        plt.figure(figsize=(10, 10))
        plt.imshow(img)
        
        # 绘制检测框
        if len(bboxes) > 0:
            for bbox, label in zip(bboxes, labels):
                if bbox[4] >= args.score_thr:
                    x1, y1, x2, y2 = bbox[:4].astype(np.int32)
                    plt.gca().add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, 
                                                    fill=False, edgecolor='red', linewidth=2))
                    plt.text(x1, y1-5, f'{model.CLASSES[label]}: {bbox[4]:.2f}',
                            color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.7))
        
        plt.axis('off')
        plt.tight_layout()
        plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
        plt.close()
    elif input_path.is_dir():
        # 获取要处理的图片列表
        image_names = set()
        if args.image_list:
            with open(args.image_list, 'r') as f:
                image_names = {line.strip() for line in f if line.strip()}
            print(f'从文件加载了 {len(image_names)} 个图片名称')
        
        # 处理目录中的所有图片
        image_extensions = ('.jpg', '.jpeg', '.png', '.bmp')
        processed_count = 0
        for img_path in input_path.glob('*'):
            if img_path.suffix.lower() in image_extensions:
                # 如果指定了图片列表，只处理列表中的图片
                if args.image_list and img_path.name not in image_names:
                    continue
                tasks = asyncio.create_task(async_inference_detector(model, str(img_path)))
                result = await asyncio.gather(tasks)
                output_path = os.path.join(args.output_dir, f'{img_path.name}')
                
                # 读取原始图像
                img = mmcv.imread(str(img_path))
                img = mmcv.bgr2rgb(img)
                
                # 处理检测结果
                if isinstance(result[0], tuple):
                    bbox_result, segm_result = result[0]
                    if isinstance(segm_result, tuple):
                        segm_result = segm_result[0]
                else:
                    bbox_result, segm_result = result[0], None
                
                # 将检测结果转换为numpy数组
                if isinstance(bbox_result, list):
                    bboxes = []
                    labels = []
                    for i, bbox in enumerate(bbox_result):
                        if len(bbox) > 0:
                            bboxes.append(bbox)
                            labels.append(np.full((len(bbox),), i, dtype=np.int64))
                    if len(bboxes) > 0:
                        bboxes = np.concatenate(bboxes, axis=0)
                        labels = np.concatenate(labels, axis=0)
                    else:
                        bboxes = np.zeros((0, 5))
                        labels = np.zeros((0,), dtype=np.int64)
                else:
                    bboxes = bbox_result
                    labels = np.zeros((len(bbox_result),), dtype=np.int64)
                
                # 创建新的图像用于可视化
                plt.figure(figsize=(10, 10))
                plt.imshow(img)
                
                # 绘制检测框
                if len(bboxes) > 0:
                    for bbox, label in zip(bboxes, labels):
                        if bbox[4] >= args.score_thr:
                            x1, y1, x2, y2 = bbox[:4].astype(np.int32)
                            plt.gca().add_patch(plt.Rectangle((x1, y1), x2-x1, y2-y1, 
                                                            fill=False, edgecolor='red', linewidth=2))
                            plt.text(x1, y1-5, f'{model.CLASSES[label]}: {bbox[4]:.2f}',
                                    color='red', fontsize=12, bbox=dict(facecolor='white', alpha=0.7))
                
                plt.axis('off')
                plt.tight_layout()
                plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
                plt.close()
                processed_count += 1
        
        print(f'总共处理了 {processed_count} 张图片')
    else:
        raise ValueError(f'输入路径不存在: {args.img}')


if __name__ == '__main__':
    args = parse_args()
    if args.async_test:
        asyncio.run(async_main(args))
    else:
        main(args)
